;;;; control variance

(load "./monte-carlo.ss")
(load "./random-number-generator.ss")


 ;;;; 1. common and antithetic random variables
;;If you want to estimate E[X-Y], X,Y are independent random-variables,
;;To get E[X-Y], you do not need the independency of X,Y;
;;Generate X,Y by same random U of uniform-dist


;;;; 2. control variable
;;control variable method seems not good enough
;;E[X] = E[X - alpha * (C-r)], where C is known random variable with mean r.
;; choose some alpha can make rhs having lower variance.

(define (control-variable H sample-strm C r)
  (let ([H-strm (stream-map H sample-strm)]
        [C-strm (stream-map C sample-strm)])
    (let ([mean-H-strm (mean H-strm)]
          [mean-HC-strm (mean (stream-map * H-strm C-strm))]
          [mean-C2-strm (mean (stream-map (lambda (ele) (expt ele 2))
                                          C-strm))])
      (let* ([alpha-strm (stream-map (lambda (mhc mh mc2)
                                       (/ (- mhc (* r mh))
                                          (- mc2 (expt r 2))))
                                     mean-HC-strm
                                     mean-H-strm
                                     mean-C2-strm)]
             [new-H-strm (stream-map (lambda (h c alpha)
                                       (- h (* alpha (- c r))))
                                     H-strm
                                     C-strm
                                     alpha-strm)])
        (mean new-H-strm))))
  )

#| ;;example: estimate Pi
(let ([H (lambda (x) (* 4 (sqrt (- 1 (expt x 2)))))]
      [C (lambda (x) (- 1 x))]
      [samples (sample-stream (lambda () (random 1.0)))]
      [r 0.5]
      [n 1000])
  (let ([old (mean (stream-map H samples))]
        [cv (control-variable H samples C r)])
    (/ (abs (- 3.1415936 (stream-ref old n)))
       (abs (- 3.1415936 (stream-ref cv n))))))
|#


;;;;; 3. stratification

;;How to stratify a sampling-flow is a hard work or depends on the problem,
;;but what to do after knowing how to stratify can be abstracted.

;;Suppose the problem has been stratified into K subproblems:
;; (p1 mc-strm1), .., (pK mc-strmK); mc-strm = stream of estimate=(mu . err) 
;; note: sum of pj = 1 with pj >0.

;;IDEA: control each level of strm at same errors, i.e. p1*err1 ~ p2*err2 ~ ..

(define (stratify-monte-carlo H . pss)
  ;;pss = (list (p1 . sampler1) ..)
  ;;return: mu streams
  (assert (pair? (car pss)))
  ;;pre-calculate at least 10000 times by probs
  (let* ([probs (map car pss)]
         [minp (apply min probs)]
         [prenums (map (lambda (p) (exact (round
                                           (* 1000 (/ p minp)))))
                       probs)]
         [estimate-streams ;(list strm1 ..), a strm: stream of (mu . sigma)
          (map (lambda (ps n0)
                 (stream-drop n0 (monte-carlo H (cdr ps) 'sigma)))
               pss
               prenums)])
    ;;calculate each n
    (let* ([psigmas (map (lambda (p esti)
                           (* p (cdr (stream-car esti))))
                         probs
                         estimate-streams)]
           [min-psig (apply min psigmas)]
           [ns (map (lambda (psig)
                      (exact (round (* 100 (/ psig min-psig)))))
                    psigmas)])
      (let loop ([es estimate-streams])
        (stream-cons (apply + (map (lambda (p esti) (* p (car (stream-car esti))))
                                   probs es))
                     (loop (map (lambda (n s) (stream-drop n s))
                                ns es)))))
    ))

#|

;;example
;;1-dim Integration in [0,1]
;;split [0,1] into K parts: jth x in [j/K, (j+1)/K], pj=1/K.

|#
(import (range))
(define (stratify-integral func K)
  (let* ([omegas (map (lambda (j)
                        (cons (/ j K) (/ (+ j 1) K)))
                      (range K))]
         [samplers (map (lambda (o) (cons (/ 1. K) (uniform (* 1. (car o)) (* 1. (cdr o)))))
                        omegas)])
    (apply stratify-monte-carlo (cons func samplers))))
;;test
;;Integrate h(x)=2*x with x in [0,1];
(let* ([h (lambda (x) (* 2 x))]
       [ordinary-version (integral-with-errors h 1)]
       [stratify-version (stratify-integral h 200)])
  (display (time (abs (- 1.0 (car (stream-ref ordinary-version 300000))))))
  ;;2e-4
  (time (abs (- 1.0 (stream-ref stratify-version 5))))
  ;;1e-6
  )



;;;; 4. importance sampling
;; estimate Pi
;; use func H(x) = 4 * sqrt(1-x^2)
(let* ([H (lambda (x) (* 4 (sqrt (- 1 (* x x)))))]
       [Pi-stream (integral-with-errors H 1)]
       [g (lambda (x) (* 2 (- 1 x)))]
       [inverseG (lambda (u) (- 1 (sqrt (- 1 u))))]
       [newH (lambda (x) (/ (H x) (g x)))]
       [Pi-stream-importance (monte-carlo newH (lambda () (inverseG (random 1.0))) 'absolute)]
       [N 100000])
  (list (stream-ref Pi-stream N)
        (stream-ref Pi-stream-importance N))
  )

